import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt

# Value to change
num_tests = 20
source_success = 0
target_success = 0
epsilon = 2.0

np.random.seed(42)
torch.manual_seed(42)

num_samples = 100
class0 = np.random.randn(num_samples, 2) * 0.5 + np.array([0, 0])
class1 = np.random.randn(num_samples, 2) * 0.5 + np.array([2, 2])
X = np.vstack((class0, class1))
y = np.hstack((np.zeros(num_samples), np.ones(num_samples)))

X_tensor = torch.from_numpy(X).float()
y_tensor = torch.from_numpy(y).long()

class SimpleMLP(nn.Module):
    def __init__(self):
        super(SimpleMLP, self).__init__()
        self.fc1 = nn.Linear(2, 10)
        self.fc2 = nn.Linear(10, 2)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

def train_model(model, X, y, epochs=100):
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.01)
    for _ in range(epochs):
        optimizer.zero_grad()
        outputs = model(X)
        loss = criterion(outputs, y)
        loss.backward()
        optimizer.step()
    return model

source_model = SimpleMLP()
source_model = train_model(source_model, X_tensor, y_tensor)

torch.manual_seed(43)
target_model = SimpleMLP()
target_model = train_model(target_model, X_tensor, y_tensor)

def fgsm_attack(data, epsilon, data_grad):
    sign_data_grad = data_grad.sign()
    perturbed_data = data + epsilon * sign_data_grad
    return perturbed_data

def plot_decision_boundary(model, X, y, title, original=None, perturbed=None):
    x_min, x_max = X[:, 0].min() - 1, X[:, 0].max() + 1
    y_min, y_max = X[:, 1].min() - 1, X[:, 1].max() + 1
    xx, yy = np.meshgrid(np.arange(x_min, x_max, 0.1), np.arange(y_min, y_max, 0.1))
    grid = torch.from_numpy(np.c_[xx.ravel(), yy.ravel()]).float()
    Z = model(grid).detach().numpy()[:, 1] > 0.5  # Probability for class 1
    Z = Z.reshape(xx.shape)
    plt.contourf(xx, yy, Z, alpha=0.8, cmap='coolwarm')
    plt.scatter(X[:, 0], X[:, 1], c=y, edgecolors='k', cmap='coolwarm')
    if original is not None:
        plt.scatter(original[0], original[1], c='green', marker='*', s=200, label='Original')
    if perturbed is not None:
        plt.scatter(perturbed[0], perturbed[1], c='orange', marker='^', s=200, label='Perturbed')
    plt.title(title)
    plt.legend()
    plt.savefig(f'{title.lower().replace(" ", "_")}.png')
    plt.show()

for i in range(num_tests):
    sample_np = np.random.randn(1, 2) * 0.1 + np.array([[0.5, 0.5]])
    sample = torch.from_numpy(sample_np).float()
    sample.requires_grad = True

    source_output = source_model(sample)
    source_pred = torch.argmax(source_output, dim=1).item()

    if source_pred != 0: continue

    loss = nn.CrossEntropyLoss()(source_output, torch.tensor([0]))
    source_model.zero_grad()
    loss.backward()
    data_grad = sample.grad.data

    perturbed_sample = fgsm_attack(sample, epsilon, data_grad)

    source_pert_output = source_model(perturbed_sample)
    source_pert_pred = torch.argmax(source_pert_output, dim=1).item()
    if source_pert_pred != source_pred:
        source_success += 1

    target_output = target_model(sample)
    target_pred = torch.argmax(target_output, dim=1).item()
    target_pert_output = target_model(perturbed_sample)
    target_pert_pred = torch.argmax(target_pert_output, dim=1).item()
    if target_pert_pred != target_pred:
        target_success += 1

    if i == 0:
        print(f'Source original pred: {source_pred}')
        print(f'Source perturbed pred: {source_pert_pred}')
        print(f'Target original pred: {target_pred}')
        print(f'Target perturbed pred: {target_pert_pred}')
        print('Original sample:', sample.detach().numpy())
        print('Perturbed sample:', perturbed_sample.detach().numpy())

        plot_decision_boundary(source_model, X, y, 'Source Model Decision Boundary', sample_np[0], perturbed_sample.detach().numpy()[0])
        plot_decision_boundary(target_model, X, y, 'Target Model Decision Boundary', sample_np[0], perturbed_sample.detach().numpy()[0])

print(f'\nAttack Success Rate (over {num_tests} samples):')
print(f'Source: {source_success / num_tests * 100:.2f}%')
print(f'Target (Transfer): {target_success / num_tests * 100:.2f}%')